#!/usr/bin/env python
# coding: utf-8

# # Model基本使用
# 通常情况下，定义训练和评估网络并直接运行，已经可以满足基本需求。
# 
# 一方面，`Model`可以在一定程度上简化代码。例如：无需手动遍历数据集；在不需要自定义`nn.TrainOneStepCell`的场景下，可以借助`Model`自动构建训练网络；可以使用`Model`的`eval`接口进行模型评估，直接输出评估结果，无需手动调用评价指标的`clear`、`update`、`eval`函数等。
# 
# 另一方面，`Model`提供了很多高阶功能，如数据下沉、混合精度等，在不借助`Model`的情况下，使用这些功能需要花费较多的时间仿照`Model`进行自定义。
# 
# 本文档首先对luojianet_ms的Model进行基本介绍，然后重点讲解如何使用`Model`进行模型训练、评估和推理。
# 
# ![model](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r1.7/tutorials/source_zh_cn/advanced/train/images/model.png)
# 
# ## Model基本介绍
# 
# [Model](http://58.48.42.237/luojiaNet/luojiaNetapi/ #luojianet_ms.Model)是luojianet_ms提供的高阶API，可以进行模型训练、评估和推理。其接口的常用参数如下：
# 
# - `network`：用于训练或推理的神经网络。
# - `loss_fn`：所使用的损失函数。
# - `optimizer`：所使用的优化器。
# - `metrics`：用于模型评估的评价函数。
# - `eval_network`：模型评估所使用的网络，未定义情况下，`Model`会使用`network`和`loss_fn`进行封装。
# 
# `Model`提供了以下接口用于模型训练、评估和推理：
# 
# - `train`：用于在训练集上进行模型训练。
# - `eval`：用于在验证集上进行模型评估。
# - `predict`：用于对输入的一组数据进行推理，输出预测结果。
# 
# ### 使用Model接口
# 
# 对于简单场景的神经网络，可以在定义`Model`时指定前向网络`network`、损失函数`loss_fn`、优化器`optimizer`和评价函数`metrics`。
# 
# 此时，`Model`会使用`network`作为前向网络，并使用`nn.WithLossCell`和`nn.TrainOneStepCell`构建训练网络，使用`nn.WithEvalCell`构建评估网络。

# In[1中低阶API实现深度学习]:


import numpy as np
import luojianet_ms.dataset as ds
import luojianet_ms.nn as nn
from luojianet_ms import Model
from luojianet_ms.common.initializer import Normal

def get_data(num, w=2.0, b=3.0):
    """生成样本数据及对应的标签"""
    for _ in range(num):
        x = np.random.uniform(-10.0, 10.0)
        noise = np.random.normal(0, 1)
        y = x * w + b + noise
        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)

def create_dataset(num_data, batch_size=16):
    """生成数据集"""
    dataset = ds.GeneratorDataset(list(get_data(num_data)), column_names=['data', 'label'])
    dataset = dataset.batch(batch_size)
    return dataset

class LinearNet(nn.Module):
    """定义线性回归网络"""
    def __init__(self):
        super().__init__()
        self.fc = nn.Dense(1, 1, Normal(0.02), Normal(0.02))

    def forward(self, x):
        return self.fc(x)

train_dataset = create_dataset(num_data=160)
net = LinearNet()
crit = nn.MSELoss()
opt = nn.Momentum(net.trainable_params(), learning_rate=0.005, momentum=0.9)

# 使用Model构建训练网络
model = Model(network=net, loss_fn=crit, optimizer=opt, metrics={"mae"})


# ### 模型训练
# 
# 使用`train`接口执行模型训练，`train`接口的常用参数如下：
# 
# - `epoch`：训练执行轮次，通常每个epoch都会使用全量数据集进行训练。
# - `train_dataset`：一个训练数据集迭代器。
# - `callbacks`：训练过程中需要执行的回调对象或者回调对象列表。
# 
# 有意思的是，如果网络模型定义了`loss_fn`，则数据和标签会被分别传给`network`和`loss_fn`，此时数据集需要返回一个元组（data, label）。如果数据集中有多个数据或者标签，可以设置`loss_fn`为None，并在`network`中实现自定义损失函数，此时数据集返回的所有数据组成的元组（data1, data2, data3, …）会传给`network`。
# 
# 如下示例使用`train`接口执行模型训练，通过`LossMonitor`回调函数查看在训练过程中的损失函数值。

# In[2高级数据集管理]:

from luojianet_ms.train.callback import LossMonitor


model.train(1, train_dataset, callbacks=[LossMonitor()])


# ### 模型评估
# 
# 使用`eval`接口进行评估，`eval`接口参数如下：
# 
# - `valid_dataset`：评估模型的数据集。
# - `callbacks`：评估过程中需要执行的回调对象或回调对象列表。
# - `dataset_sink_mode`：数据是否直接下沉至处理器进行处理。

# In[3图像处理]:


eval_dataset = create_dataset(num_data=80)  # 创建评估数据集
eval_result = model.eval(eval_dataset)      # 执行模型评估
print(eval_result)


# ### 模型推理
# 
# 使用`predict`接口进行推理，`predict`接口参数如下：
# 
# - `predict_data`：预测样本，数据可以是单个张量、张量列表或张量元组。

# In[4自然语言]:


eval_data = eval_dataset.create_dict_iterator()
data = next(eval_data)
# 执行模型预测
output = model.predict(data["data"])
print(output)


# 一般情况下需要对推理结果进行后处理才能得到比较直观的推理结果。
# 
# ## 自定义场景
# 
# luojianet_ms提供的网络封装函数`nn.WithLossCell`、`nn.TrainOneStepCell`和`nn.WithEvalCell`并不适用于所有场景，实际场景中常常需要自定义网络的封装函数，这种情况下`Model`使用这些封装函数自动地进行封装显然是不合理的。
# 
# 接下来介绍在自定义网络封装函数时如何正确地使用`Model`。

# ### 自定义损失网络
# 
# 在有多个数据或者多个标签的场景下，可以使用自定义损失网络将前向网络和自定义的损失函数链接起来作为`Model`的`network`，`loss_fn`使用默认值`None`，此时`Model`内部不会经过`nn.WithLossCell`，而会直接使用`nn.TrainOneStepCell`将`network`与`optimizer`组成训练网络。

# In[5]:


import numpy as np
import luojianet_ms.dataset as ds
import luojianet_ms.ops as ops
import luojianet_ms.nn as nn
from luojianet_ms import Model
from luojianet_ms.nn import LossBase
from luojianet_ms.train.callback import LossMonitor

def get_multilabel_data(num, w=2.0, b=3.0):
    """生成多标签数据，产生一组数据x对应两个标签y1和y2"""
    for _ in range(num):
        x = np.random.uniform(-10.0, 10.0)
        noise1 = np.random.normal(0, 1)
        noise2 = np.random.normal(-1, 1)
        y1 = x * w + b + noise1
        y2 = x * w + b + noise2
        yield np.array([x]).astype(np.float32), np.array([y1]).astype(np.float32), np.array([y2]).astype(np.float32)

def create_multilabel_dataset(num_data, batch_size=16):
    """生成多标签数据集，一个数据data对应两个标签label1和label2"""
    dataset = ds.GeneratorDataset(list(get_multilabel_data(num_data)), column_names=['data', 'label1', 'label2'])
    dataset = dataset.batch(batch_size)
    return dataset

class L1LossForMultiLabel(LossBase):
    """自定义多标签损失函数"""

    def __init__(self, reduction="mean"):
        super(L1LossForMultiLabel, self).__init__(reduction)
        self.abs = ops.Abs()

    def forward(self, base, target1, target2):
        """输入有三个，分别为预测值base，真实值target1和target2"""
        x1 = self.abs(base - target1)
        x2 = self.abs(base - target2)
        return self.get_loss(x1) / 2 + self.get_loss(x2) / 2

class CustomWithLossCell(nn.Module):
    """连接前向网络和损失函数"""

    def __init__(self, backbone, loss_fn):
        """输入有两个，前向网络backbone和损失函数loss_fn"""
        super(CustomWithLossCell, self).__init__(auto_prefix=False)
        self._backbone = backbone
        self._loss_fn = loss_fn

    def forward(self, data, label1, label2):
        output = self._backbone(data)                 # 前向计算得到网络输出
        return self._loss_fn(output, label1, label2)  # 得到多标签损失值

multi_train_dataset = create_multilabel_dataset(num_data=160)

# 构建线性回归网络
net = LinearNet()
# 多标签损失函数
loss = L1LossForMultiLabel()

# 连接线性回归网络和多标签损失函数
loss_net = CustomWithLossCell(net, loss)
opt = nn.Momentum(net.trainable_params(), learning_rate=0.005, momentum=0.9)

# 使用Model连接网络和优化器，此时Model内部不经过nn.WithLossCell
model = Model(network=loss_net, optimizer=opt)
# 使用train接口进行模型训练
model.train(epoch=1, train_dataset=multi_train_dataset, callbacks=[LossMonitor()])


# ### 自定义训练网络
# 
# 在自定义训练网络时，需要手动构建训练网络作为`Model`的`network`，`loss_fn`和`optimizer`均使用默认值`None`，此时`Model`会使用`network`作为训练网络，而不会进行任何封装。
# 
# 如下示例自定义训练网络`CustomTrainOneStepCell`，然后通过`Model`接口构建训练网络。

# In[6]:


import luojianet_ms.ops as ops
from luojianet_ms import Model
from luojianet_ms.train.callback import LossMonitor

class CustomTrainOneStepCell(nn.Module):
    """自定义训练网络"""

    def __init__(self, network, optimizer, sens=1.0):
        """入参有三个：训练网络，优化器和反向传播缩放比例"""
        super(CustomTrainOneStepCell, self).__init__(auto_prefix=False)
        self.network = network                    # 定义前向网络
        self.network.set_grad()                   # 构建反向网络
        self.optimizer = optimizer                # 定义优化器
        self.weights = self.optimizer.parameters  # 待更新参数
        self.grad = ops.GradOperation(get_by_list=True, sens_param=True)  # 反向传播获取梯度

    def forward(self, *inputs):
        loss = self.network(*inputs)                    # 执行前向网络，计算当前输入的损失函数值
        grads = self.grad(self.network, self.weights)(*inputs, loss)  # 进行反向传播，计算梯度
        loss = ops.depend(loss, self.optimizer(grads))  # 使用优化器更新梯度
        return loss

multi_train_ds = create_multilabel_dataset(num_data=160)

# 手动构建训练网络
train_net = CustomTrainOneStepCell(loss_net, opt)
# 构建训练网络
model = Model(train_net)
# 执行模型训练
model.train(epoch=1, train_dataset=multi_train_ds, callbacks=[LossMonitor()])


# ### 自定义评估网络
# 
# `Model`默认使用`nn.WithEvalCell`构建评估网络，在不满足需求的情况下需要手动构建评估网络，如多数据和多标签场景下。
# 
# 如下示例自定义评估网络`CustomWithEvalCell`，然后使用`Model`接口构建评估网络。

# In[7]:


import luojianet_ms.nn as nn
from luojianet_ms.train.callback import LossMonitor
from luojianet_ms import Model


class CustomWithEvalCell(nn.Module):
    """自定义多标签评估网络"""

    def __init__(self, network):
        super(CustomWithEvalCell, self).__init__(auto_prefix=False)
        self.network = network

    def forward(self, data, label1, label2):
        output = self.network(data)
        return output, label1, label2

# 构建多标签评估数据集
multi_eval_dataset = create_multilabel_dataset(num_data=80)

# 构建评估网络
eval_net = CustomWithEvalCell(net)

# 评估函数
mae1 = nn.MAE()
mae2 = nn.MAE()
mae1.set_indexes([0, 1])
mae2.set_indexes([0, 2])

# 使用Model构建评估网络
model = Model(network=loss_net, optimizer=opt, eval_network=eval_net,
              metrics={"mae1": mae1, "mae2": mae2})
result = model.eval(multi_eval_dataset)
print(result)


# 上述代码在进行模型评估时，评估网络的输出会透传给评估指标的`update`函数，`update`函数将接收到三个输入，分别为`logits`、`label1`和`label2`。
# 
# `nn.MAE`仅允许在两个输入上计算评价指标，因此使用`set_indexes`指定`mae1`使用下标为0和1的输入，也就是`logits`和`label1`，计算评估结果；指定`mae2`使用下标为0和2的输入，也就是`logits`和`label2`，计算评估结果。

# ### 网络推理
# 
#    `Model`没有提供用于指定自定义推理网络的参数，此时可以直接运行前向网络获得推理结果。

# In[8]:


for d in multi_eval_dataset.create_dict_iterator():
    data = d["data"]
    break

output = net(data)
print(output)

